from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
import random
import sys
import os

# 添加项目根目录到 Python 路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.database import SQLALCHEMY_DATABASE_URL

def update_student_grades():
    """为所有学生随机设置年级（1-4）"""
    # 创建数据库连接
    engine = create_engine(SQLALCHEMY_DATABASE_URL)
    Session = sessionmaker(bind=engine)
    db = Session()

    try:
        # 获取所有学生ID
        result = db.execute(text("SELECT id FROM students"))
        student_ids = [row[0] for row in result]
        
        print(f"开始为 {len(student_ids)} 名学生更新年级信息...")
        
        # 为每个学生设置随机年级
        for student_id in student_ids:
            grade = random.randint(1, 4)
            db.execute(
                text("UPDATE students SET grade = :grade WHERE id = :id"),
                {"grade": grade, "id": student_id}
            )
        
        # 提交更改
        db.commit()
        print("年级信息更新完成！")
        
        # 统计年级分布
        result = db.execute(text("""
            SELECT grade, COUNT(*) as count 
            FROM students 
            GROUP BY grade 
            ORDER BY grade
        """))
        stats = result.fetchall()
        
        print("\n年级分布统计：")
        total = 0
        for grade, count in stats:
            if grade is not None:  # 只显示有年级的学生
                print(f"大{grade}：{count}人")
                total += count
        print(f"\n总计：{total}人")

        # 验证是否还有学生没有年级信息
        result = db.execute(text("SELECT COUNT(*) FROM students WHERE grade IS NULL"))
        null_count = result.scalar()
        if null_count > 0:
            print(f"\n警告：仍有 {null_count} 名学生没有年级信息！")
        else:
            print("\n验证完成：所有学生都已设置年级信息！")

    except Exception as e:
        print(f"更新过程中发生错误: {str(e)}")
        db.rollback()
    finally:
        db.close()

if __name__ == "__main__":
    update_student_grades() 